/*
 * Copyright 2024 ywcai
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

`include "defines.v"
`timescale 1ns/1ps

module trap(
    input	wire					clk,
    input	wire					rst_n,

    input	wire[`MemAddrBus]		pc_i,
    input	wire					mie_i,
    input	wire[`RegDataBus]		csr_mtvec_i,
    input	wire[`RegDataBus]		csr_mepc_i,
    input	wire[`RegDataBus]		csr_mstatus_i,
	input	wire[`TrapBus]			trap_code_i,

    output	reg						csr_we_o,
    output	reg[`MemAddrBus]		csr_wr_addr_o,
    output	reg[`RegDataBus]		csr_wr_data_o,
    output	reg[`MemAddrBus]		trap_entry_o,
    output	reg						trap_assert_o,
    output	wire					flush_id_ex_req_o
    );

    localparam S_TRAP_IDLE			= 4'b0001;
    localparam S_TRAP_SYNC_ASSERT	= 4'b0010;
    localparam S_TRAP_ASYNC_ASSERT	= 4'b0100;
    localparam S_TRAP_MRET			= 4'b1000;

    localparam S_CSR_IDLE			= 5'b00001;
    localparam S_CSR_MSTATUS		= 5'b00010;
    localparam S_CSR_MEPC			= 5'b00100;
    localparam S_CSR_MSTATUS_MRET	= 5'b01000;
    localparam S_CSR_MCAUSE			= 5'b10000;

    reg[3:0]	trap_state;
    reg[4:0]	csr_state, csr_next_state;
    reg[`MemAddrBus]	inst_addr;
    reg[`RegDataBus]	cause;
    wire[`IrqBus]	irqs		= trap_code_i[7:0];
	wire[`ExcBus]	exceptions	= trap_code_i[15:8];
	wire	e_inst_illegal		= trap_code_i[8];
	wire	e_inst_ecall		= trap_code_i[9];
	wire	e_inst_ebreak		= trap_code_i[10];
	wire	e_inst_mret			= trap_code_i[12];
	wire	e_inst_misalign		= trap_code_i[13];
	wire	e_data_misalign		= trap_code_i[14];

	assign flush_id_ex_req_o = ((trap_state != S_TRAP_IDLE)
		| (csr_state != S_CSR_IDLE)) ? `ENABLE: `DISABLE;

    always @(*) begin
		if (rst_n == `RESET_ENABLE) begin
			trap_state = S_TRAP_IDLE;
		end else begin
			if ((exceptions != `EXC_NONE) & !e_inst_mret) begin
				trap_state = S_TRAP_SYNC_ASSERT;
			end else if ((irqs != `IRQ_NONE) & (mie_i == `ENABLE)) begin
				trap_state = S_TRAP_ASYNC_ASSERT;
			end else if (e_inst_mret) begin
				trap_state = S_TRAP_MRET;
			end else begin
				trap_state = S_TRAP_IDLE;
			end
		end
	end

    always @(posedge clk) begin
		if (rst_n == `RESET_ENABLE) begin
			csr_state <= S_CSR_IDLE;
		end else begin
			csr_state <= csr_next_state;
		end
	end

	always @(*) begin
		if (rst_n == `RESET_ENABLE) begin
			csr_next_state = S_CSR_IDLE;
		end else begin
			case (csr_state)
				S_CSR_IDLE: begin
					if (trap_state == S_TRAP_SYNC_ASSERT) begin
						csr_next_state = S_CSR_MEPC;
						inst_addr = pc_i;

						if (e_inst_illegal) begin
							cause = `XLEN'd2;
						end else if (e_inst_ebreak) begin
							cause = `XLEN'd3;
						end else if (e_inst_ecall) begin
							cause = `XLEN'd11;
						end else if (e_inst_misalign) begin
							cause = `XLEN'd0;
						end else if (e_data_misalign) begin
							cause = `XLEN'd6;
						end else begin
							cause = `XLEN'hff;
						end
					end else if (trap_state == S_TRAP_ASYNC_ASSERT) begin
						cause = `XLEN'h80000004;
						csr_next_state = S_CSR_MEPC;
						inst_addr = pc_i;
					end else if (trap_state == S_TRAP_MRET) begin
						csr_next_state = S_CSR_MSTATUS_MRET;
					end
				end

				S_CSR_MEPC: begin
					csr_next_state = S_CSR_MSTATUS;
				end

				S_CSR_MSTATUS: begin
					csr_next_state = S_CSR_MCAUSE;
				end

				S_CSR_MCAUSE: begin
					csr_next_state = S_CSR_IDLE;
				end

				S_CSR_MSTATUS_MRET: begin
					csr_next_state = S_CSR_IDLE;
				end

				default: begin
					csr_next_state = S_CSR_IDLE;
				end
			endcase
		end
    end

	always @(posedge clk) begin
		if (rst_n == `RESET_ENABLE) begin
			csr_we_o <= `DISABLE;
			csr_wr_addr_o <= `ZERO_ADDR;
			csr_wr_data_o <= `ZERO;
		end else begin
			case (csr_state)
				S_CSR_MEPC: begin
					csr_we_o <= `ENABLE;
					csr_wr_addr_o <= {28'h0, `CSR_MEPC};
					csr_wr_data_o <= inst_addr;
				end

				S_CSR_MCAUSE: begin
					csr_we_o <= `ENABLE;
					csr_wr_addr_o <= {28'h0, `CSR_MCAUSE};
					csr_wr_data_o <= cause;
				end

				S_CSR_MSTATUS: begin
					csr_we_o <= `ENABLE;
					csr_wr_addr_o <= {28'h0, `CSR_MSTATUS};
					csr_wr_data_o <= {32'h0, csr_mstatus_i[31:4], 1'b0, csr_mstatus_i[2:0]};
				end

				S_CSR_MSTATUS_MRET: begin
					csr_we_o <= `ENABLE;
					csr_wr_addr_o <= {28'h0, `CSR_MSTATUS};
					csr_wr_data_o <= {32'h0, csr_mstatus_i[31:4], csr_mstatus_i[7], csr_mstatus_i[2:0]};
				end

				default: begin

				end
			endcase
		end
	end

	always @(posedge clk) begin
		if (rst_n == `RESET_ENABLE) begin
			trap_assert_o <= `IRQ_DEASSERT;
			trap_entry_o <= `ZERO_ADDR;
		end else begin
			case (csr_state)
				S_CSR_MCAUSE: begin
					trap_assert_o <= `IRQ_ASSERT;
					trap_entry_o <= csr_mtvec_i[`ADDR_MSB:0];
				end

				S_CSR_MSTATUS_MRET: begin
					trap_assert_o <= `IRQ_ASSERT;
					trap_entry_o <= csr_mepc_i[`ADDR_MSB:0];
				end

				default: begin

				end
			endcase
		end
	end

endmodule
